package com.yiquan.video.config;

import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;


public class MdcThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

    @Override
    public void execute(Runnable task) {

        Map currentMDC = MDC.getCopyOfContextMap();
        super.execute(() -> {
            MDC.setContextMap(currentMDC);
            try {
                task.run();
            } finally {
                MDC.clear();
            }
        });
    }

    @Override
    public void execute(Runnable task, long startTimeout) {
        Map currentMDC = MDC.getCopyOfContextMap();
        super.execute(() -> {
            MDC.setContextMap(currentMDC);
            try {
                task.run();
            } finally {
                MDC.clear();
            }
        }, startTimeout);
    }


    @Override
    public <T> Future<T> submit(Callable<T> task) {

        Map previous = MDC.getCopyOfContextMap();

        return super.submit(() -> {
            MDC.setContextMap(previous);
            try {
                return task.call();
            } finally {
                MDC.clear();
            }
        });

    }


}